# 把标注数据转变成yolo格式 并分割数据 并开始跑训练 by LR
import json2txt as j2t
import dividata as di
import os
import shutil
from ultralytics import YOLO

def delete_files_in_folder(folder_path):

    if not os.path.exists(folder_path):
        return

    for filename in os.listdir(folder_path):
        file_path = os.path.join(folder_path, filename)
        try:
            if os.path.isfile(file_path) or os.path.islink(file_path):
                os.unlink(file_path)
            elif os.path.isdir(file_path):
                shutil.rmtree(file_path)
        except Exception as e:
            print(f'Failed to delete {file_path}. Reason: {e}')

if __name__=='__main__':
    #标注数据文件夹
    srcDir = r"E:\LR\handian20240801\biaozhu"
    #标注文件复制出来的图片和生成txt 的路径
    img_dir = r"E:\LR\handian20240801\imgs"
    # 类别和序号的映射字典
    cls2id_dict = {"bodian": "0", "wubodian": "1", "handian": "2", "konghan": "3", "xuhan": "4", "xizhu": "5"}

    #随机分割训练集测试集 图片路径
    trainDir = r'E:\LR\handian20240801\train'
    valDir = r'E:\LR\handian20240801\val'
    # 训练集测试集分割比例
    radio = 0.8

    #训练配置
    yaml = 'handian20240718.yaml'

    #删除文件重新建立
    if not os.path.exists(img_dir):
        os.makedirs(img_dir)
    delete_files_in_folder(img_dir)

    # 生成图片和生成txt
    list = os.listdir(srcDir)
    for f in list:
        if f.endswith('.jpg'):
            img_src = os.path.join(srcDir, f)
            shutil.copy(img_src, os.path.join(img_dir, f))
        if f.endswith('.json'):
            json_path = os.path.join(srcDir, f)
            txt_name = f.split(".")[0] + ".txt"
            save_txt_path = os.path.join(img_dir, txt_name)
            labels = j2t.labelme_to_yolo(json_path, cls2id_dict)
            j2t.write_label2txt(save_txt_path, labels)

    # 先清空文件夹中内容
    delete_files_in_folder(trainDir)
    delete_files_in_folder(valDir)

    # 分割训练集和测试集
    di.split_dataset(img_dir,trainDir,valDir,radio)

    #开始训练
    model = YOLO("yolov8n.pt")  # load a pretrained model (recommended for training)

    # Use the model
    model.train(data=yaml, epochs=200,imgsz=1024,batch=12,device=0)  # train the model
    model.export(format="onnx")